#!/usr/bin/env python
# -*- coding: utf-8 -*-
#    Copyright (C) 2012  Ruslan V. Us <unclerus@gmail.com>, Aleksandr V. Bannov <kokudzin@gmail.com>
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

import time
import _daemon
import _client
import _providers
import _settings
import argparse
import subprocess
import threading
from string import strip

clientsLock = threading.Lock ()

class Daemon (_daemon.BaseDaemon):
    
    def __init__ (self, *args, **kwargs):
        super (Daemon, self).__init__ (*args, **kwargs)
        self.clients = {}
        self.banned = []
        self.providers = []
        self.notBan = kwargs.get ('n')
        self.configSection = kwargs.get ('s')

        self.configFile = kwargs.get ('c')
        self._whiteList = None
    
    def addProvider (self, provider):
                  
        if provider not in self.providers:
            self.providers.append (provider)
            
    def loadWhiteList (self, fileName):
        print 'Loading white list...'
        try:    
            whiteListFile = file (fileName, 'r')
        except Exception as err:
            print 'error.'
            self._log.error (str (err), 'loadWhiteList error', 50, False)  
            return False
        
        self._whiteList = map (strip, whiteListFile.readlines ())
        print 'done.' 
        
    def loadFilters (self, fileName):
        print 'Loading filters...'
        try:    
            filters = file (fileName, 'r')
        except Exception as err:
            print 'error.'
            self._log.error (str (err), 'loadFilters error', 50, False)  
            return False
        self.filters = []

        self.filters = map (strip, filters.readlines ())
        print 'done.' 
        
    def isWhite (self, ip):
        if ip in self._whiteList:
            return True
        return False
    
    def banClient (self, client):
        try:
            if self.notBan == False:  
                subprocess.call (
                    '/sbin/iptables -A INPUT -p tcp --dport {port} -s {ip} -j DROP'.format (port = '80', ip = client.ip).split ()
                )      
                self._log.notice ('{ip} banned'.format (ip = client.ip))
                self._log.notice (str (client))
                self._log.notice ('\n'.join(value for value in client.stat))
                self.banned.append (client)
        except  Exception as e:
            self._log.error (str (e), 'banClient error', 50, False)  
    
    def processRequest (self, request):
        if self._whiteList != None and self.isWhite (request.ip):
            return
        with clientsLock:
            if request.ip not in self.clients:
                self.clients [request.ip] = _client.Client (request.ip, self.settings)
            client = self.clients [request.ip]
            if client in self.banned:
                return
            client.registerActivity (request)
            if client.isBot ():
                self.banClient (client)
    
    def processProvider (self, provider):
        try:
            requests = provider.next ()
            if not requests:
                return
            for request in requests:
                self.processRequest (request)
        except StopIteration:
            pass
    
    def run (self):
        try:
            self.settings = _settings.Settings (self.configFile, self.configSection)
        except Exception as inst:
            self._log.error (str (inst), 'run error', 50, False)  
            #dTmp = Daemon ('/tmp/blocker.pid')       
            #dTmp.stop ()
            self.stop ()
        
        self.startCleaner ()            
        while True:
            for provider in self.providers:
                self.processProvider (provider)
            time.sleep (0.5)
            #self._log.notice ('\n'.join(str (self.clients [n]) for n in self.clients))
            
    def startCleaner (self):
        
        cleaner = cleanerThread (self.clients, self.settings)
        cleaner.start ()
        
class cleanerThread (threading.Thread):
    def __init__ (self, clients, settings):
        threading.Thread.__init__ (self, name = 'cleanerThread')
        self.settings = settings
        self.clients = clients
        
    def run (self):
        self.clean ()
            
    def clean (self):
        while True:
            time.sleep (self.settings.clientsCleanFrequency)
            self.clearInactiveClients ()
            
    def clearInactiveClients (self):
        with clientsLock:
            currentTime = time.time()
            unactiveClients = []
            for client in self.clients:
                if currentTime - self.clients [client].lastRequest > self.settings.clientDropTimeout:
                    unactiveClients.append (client)
            for client in unactiveClients:
                self.clients.pop (client)
                
                
options_start = (
    {
        'args': ('log',),
        'kwargs': {
            'help': 'log file to analyze',
            'metavar' : 'file.log'
        }
    },
    {
        'args': ('-n',),
        'kwargs': {
            'default' : False,
            'action' : 'store_true',
            'help': 'disable blocking',
        }        
    },
    {
        'args': ('-s',),
        'kwargs': {
            'help': 'config section',
            'metavar' : 'default',
            'default' : 'default'
        }        
    },
    {
        'args': ('-w',),
        'kwargs': {
            'help': 'white list file',
            'metavar' : 'whitelist.txt'
        }        
    },
    {
        'args': ('-c',),
        'kwargs': {
            'help': 'config file',
            'metavar' : 'defender.conf',
            'default' : 'defender.conf'
        }        
    },
    {
        'args': ('-f',),
        'kwargs': {
            'help': 'filters file',
            'metavar' : 'filters.txt'
        }
    }
)

options = (
    {
        'args': ('-p',),
        'kwargs': {
            'help': 'PID file',
            'metavar' : '/tmp/blocker.pid',
            'default' : '/tmp/blocker.pid'
        }        
    },
)

def start (args):
    d = Daemon (args.p, **args.__dict__)
    if (args.w):
        d.loadWhiteList (args.w)
    if (args.f):
        d.loadFilters (args.f)
    d.addProvider (_providers.ApacheLogProvider (args.log, d.filters) ())
    d.start ()
    
def stop (args):
    d = Daemon (args.p)
    d.stop ()
    
def restart (args):
    d = Daemon (args.p)
    d.restart ()
    
def main ():
    parser = argparse.ArgumentParser ()
    subparsers = parser.add_subparsers (help='start commands')
    
    for option in options:
        parser.add_argument (*option ['args'], **option ['kwargs'])
        
    parser_start = subparsers.add_parser ('start', help='start')
    parser_start.set_defaults (func=start)
    
    parser_stop = subparsers.add_parser ('stop', help='stop')
    parser_stop.set_defaults (func=stop)
    
    parser_restart = subparsers.add_parser ('restart', help='restart')
    parser_restart.set_defaults (func=restart)
    
    for option in options_start:
        parser_start.add_argument (*option ['args'], **option ['kwargs'])
        
    args = parser.parse_args ()
    args.func (args)

if __name__ == "__main__":
    main ()
    